# -- coding: utf-8 --
"""
能量方程求解
"""

# third party lib
import numpy as np

# private lib
from utils.constant import InterpMode
from utils.numeric import grid_interp, difference
from utils.optmization import error_SA

def cal_energy(params: dict, gas_wall_friction_force:np.ndarray, liquid_wall_friction_force:np.ndarray, gas_liquid_friction_force:np.ndarray) -> np.ndarray:
    """ 能量方程计算
    Args:
        params: dict 计算需要的输入参数
        gas_wall_friction_force: np.ndarray 气体壁面摩擦力
        liquid_wall_friction_force: np.ndarray 液体壁面摩擦力
        gas_liquid_friction_force: np.ndarray 气体液体摩擦力
        
    Returns:
        T_node_implicit_LHS: np.ndarray 温度
    """

    dx = params["dx"]
    dt = params["dt"]

    cross_section_area_segment = params["cross_section_area_segment"]
    
    T_node = params["T_node"]
    T_segment = grid_interp(T_node, InterpMode.MID)
    p_node = params["p_node"]
    p_segment = grid_interp(p_node, InterpMode.MID)
    
    rho_g_node = params["rho_g_node"]
    rho_g_segment = grid_interp(rho_g_node, InterpMode.MID)
    
    rho_L_node = params["rho_L_node"]
    rho_L_segment = grid_interp(rho_L_node, InterpMode.MID)
    
    alpha_g_node = params["alpha_g_node"]
    alpha_g_segment = grid_interp(alpha_g_node, InterpMode.MID)
    alpha_L_node = params["alpha_L_node"]
    alpha_L_segment = grid_interp(alpha_L_node, InterpMode.MID)
    
    V_g_segment = params["V_g_segment"]
    V_L_segment = params["V_L_segment"]
    
    T_node_iteration_init = params["T_node_iteration_init"]
    p_node_iteration_init = params["p_node_iteration_init"]
    p_segment_iteration_init = grid_interp(p_node_iteration_init, InterpMode.MID)
    alpha_g_node_iteration_init = params["alpha_g_node_iteration_init"]
    alpha_g_segment_iteration_init = grid_interp(alpha_g_node_iteration_init, InterpMode.MID)
    alpha_L_node_iteration_init = params["alpha_L_node_iteration_init"]
    alpha_L_segment_iteration_init = grid_interp(alpha_L_node_iteration_init, InterpMode.MID)
    rho_L_node_iteration_init = params["rho_L_node_iteration_init"]
    rho_L_segment_iteration_init = grid_interp(rho_L_node_iteration_init, InterpMode.MID)
    
    # 更新物性
    Cp_g_segment = grid_interp(params["Cp_g_node"], InterpMode.MID)
    Cp_L_segment = grid_interp(params["Cp_L_node"], InterpMode.MID)
    
    Cp_g_node_iteration_init = params["Cp_g_node_iteration_init"]
    Cp_L_node_iteration_init = params["Cp_L_node_iteration_init"]
    
    Cp_g_segment_iteration_init = grid_interp(Cp_g_node_iteration_init, InterpMode.MID)
    Cp_L_segment_iteration_init = grid_interp(Cp_L_node_iteration_init, InterpMode.MID)
    
    K = params["K"]

    # 计算焦汤系数
    # C_JT_g = - 1/Cp_g_segment/rho_g_segment*(T_segment/rho_g_segment*grid_interp(params["rho_g_derivate_p_node"], InterpMode.MID) + 1)
    # C_JT_L = - 1/Cp_L_segment/rho_L_segment*(T_segment/rho_L_segment*grid_interp(params["rho_L_derivate_p_node"], InterpMode.MID) + 1)
    C_JT_g = 0
    C_JT_L = 0

    # 初始化
    T_node_implicit_LHS = np.copy(T_node)
    T_node_implicit_RHS = np.copy(T_node)

    # 计算右端项分母    
    energy_denominator_RHS = grid_interp(
        rho_g_segment * alpha_g_segment * Cp_g_segment + rho_L_segment * alpha_L_segment * Cp_L_segment, InterpMode.MID.value
    ) / dt + K * grid_interp(params["D_segment"] * np.pi / cross_section_area_segment, InterpMode.MID.value)
    
    energy_RHS_part_1 = grid_interp(
    (alpha_g_segment*rho_g_segment*Cp_g_segment*C_JT_g + 1)*(
        p_segment * alpha_g_segment
        - p_segment_iteration_init * alpha_g_segment_iteration_init
    )

    / dt
    + (alpha_L_segment*rho_L_segment*Cp_L_segment*C_JT_L + 1)*(
        p_segment * alpha_L_segment
        - p_segment_iteration_init * alpha_L_segment_iteration_init
    )


    / dt + K * params["D_segment"] * np.pi * params["T0"] / cross_section_area_segment,
    InterpMode.MID
    )

    energy_RHS_part_2 = grid_interp(
        (alpha_g_segment*rho_g_segment*Cp_g_segment*C_JT_g + 1)*alpha_g_segment * V_g_segment * difference(p_node, dx, 1)

        + (alpha_L_segment*rho_L_segment*Cp_L_segment*C_JT_L + 1)*alpha_L_segment * V_L_segment * difference(p_node, dx, 1)

        + gas_wall_friction_force / cross_section_area_segment * V_g_segment
        + liquid_wall_friction_force / cross_section_area_segment * V_L_segment
        + gas_liquid_friction_force / cross_section_area_segment * (V_g_segment - V_L_segment),
    1,
    )
    
    error = 1
    while error > 1e-6:
        energy_RHS_part_3 = (
            -grid_interp(
                rho_g_segment
                * alpha_g_segment_iteration_init
                * Cp_g_segment_iteration_init
                * grid_interp(T_node_iteration_init, 1)
                + rho_L_segment_iteration_init
                * alpha_L_segment_iteration_init
                * Cp_L_segment_iteration_init
                * grid_interp(T_node_iteration_init, 1),
                1,
            )
            / dt
            + grid_interp(rho_g_segment * alpha_g_segment * V_g_segment * Cp_g_segment, 1) * difference(grid_interp(T_node_implicit_RHS, 1), dx, 1)
            + grid_interp(rho_L_segment * alpha_L_segment * V_L_segment * Cp_L_segment, 1) * difference(grid_interp(T_node_implicit_RHS, 1), dx, 1)
            )
        T_node_implicit_LHS[1:-1] = (
            -energy_RHS_part_3 + energy_RHS_part_1 + energy_RHS_part_2
        ) / energy_denominator_RHS

        error = error_SA(T_node_implicit_LHS, T_node_implicit_RHS)
        T_node_implicit_RHS = np.copy(T_node_implicit_LHS)

    return T_node_implicit_LHS
